from maskrcnn_benchmark.config import cfg
from predictor import COCODemo
from gevent.pywsgi import WSGIServer
from flask import Flask, request, Response, send_file
from io import BytesIO
import numpy as np
import cv2
import urllib.request
import sys

def url_to_image(url):
	# download the image, convert it to a NumPy array, and then read
	# it into OpenCV format
	resp = urllib.request.urlopen(url)
	image = np.asarray(bytearray(resp.read()), dtype="uint8")
	image = cv2.imdecode(image, cv2.IMREAD_COLOR)
 
	# return the image
	return image

# update the config options with the config file
cfg.merge_from_file("../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml")
# manual override some options

cfg.merge_from_list(["MODEL.DEVICE", "gpu"])

coco_demo = COCODemo(
    cfg,
    min_image_size=600,
    confidence_threshold=0.7,
)
# load image and then run predictio

app = Flask(__name__)


# route http posts to this method
@app.route('/api/test', methods=['POST'])
def test():
    r = request
    # decode image
    print(r.form['url'])
    img = url_to_image(r.form['url'])
    
    # do some fancy processing here....
    predictions = coco_demo.run_on_opencv_image(img)
    # build a response dict to send back to client
    """ cv2.imwrite('result.jpg', predictions)
    return 'ok' """
    _, img_encoded = cv2.imencode('.jpg', predictions)
    # send http request with image and receive response
    strIO = BytesIO()
    strIO.write(img_encoded.tobytes())
    strIO.seek(0)
    return send_file(strIO,mimetype='image/jpeg')


# start flask app
http_server = WSGIServer(('0.0.0.0', 5000), app)
http_server.serve_forever()